#ifndef PT_RRT_H
#define PT_RRT_H

#include <stdexcept>
#include <boost/shared_ptr.hpp>
#include "ompl_includes.h"
#include "StatePtr.h"
#include "FrictionModel.h"
#include "Visualizer.h"
#include "NullVisualizer.h"

class ptRRT : public ompl::base::Planner {

  const boost::shared_ptr<FrictionModel> frictionModel;
  double range;
  ompl::geometric::PathSimplifierPtr psk;
  ompl::RNG rng;
  boost::shared_ptr<Visualizer> visualizer;
  bool simplify; // simplification leads to less possible paths

public:

  ptRRT(const ompl::base::SpaceInformationPtr &si,
  const boost::shared_ptr<FrictionModel> frictionModel)
  : ompl::base::Planner(si, "ptRRT"),
    frictionModel(frictionModel),
    range(0.5),
    psk(new ompl::geometric::PathSimplifier(si)),
    visualizer(new NullVisualizer()),
    simplify(false)
  {
  }

  void setRange(double range) {
    if (range <= 0.0) {
      throw std::invalid_argument("range <= 0");
    }
    this->range = range;
  }

  void setVisualizer(boost::shared_ptr<Visualizer> v) {
    visualizer = v;
  }

  void setSimplify(bool simplify) {
    this->simplify = simplify;
  }

private:

  double estimateProbability(double x, double y) {
    double acc = 0.0;
    const int N = 13; // number of samples
    for (int i = 0; i < N; ++i) {
      acc += frictionModel->probability(x + rng.uniformReal( -0.4, 0.4 ), y + rng.uniformReal( -0.4, 0.4 ));
    }
    return acc / static_cast<double>(N);
  }

  double motionProbability(const ompl::geometric::PathGeometric & path) {
    std::size_t stateCount = path.getStateCount();
    double p = 1.0;
    for (std::size_t i = 1; i < stateCount; ++i) {
      const ompl::base::State * state = path.getState(i);
      typedef ompl::base::SE2StateSpace::StateType SE2StateType;
      const SE2StateType * se2state = dynamic_cast<const SE2StateType*>(state);
      p *= estimateProbability(se2state->getX(), se2state->getY());
    }
    return p;
  }

public:

  virtual ompl::base::PlannerStatus solve(const ompl::base::PlannerTerminationCondition &ptc) {
    checkValidity();

    std::cout << "solving with ptRRT" << std::endl;

    // get the goal
    ompl::base::GoalSampleableRegion * goalRegion = dynamic_cast<ompl::base::GoalSampleableRegion*>(pdef_->getGoal().get());
    if (!goalRegion) {
      return ompl::base::PlannerStatus::UNRECOGNIZED_GOAL_TYPE;
    }
    StatePtr goal(si_);
    goalRegion->sampleGoal(goal);

    // start state
    const ompl::base::State * start = pis_.nextStart();
    if (start == 0) {
      return ompl::base::PlannerStatus::INVALID_START;
    }

    // best solution
    double bestProbability = 0;
    ompl::base::ProblemDefinitionPtr best;
    ompl::base::PlannerStatus bestStatus;


    // start planning
    while (ptc() == false) {
      // generate a path with RRT
      ompl::base::ProblemDefinitionPtr pdef(new ompl::base::ProblemDefinition(si_));
      pdef->setStartAndGoalStates(start, goal);
      ompl::geometric::RRT rrt(si_);
      rrt.setRange(range);
      rrt.setProblemDefinition(pdef);

      ompl::base::PlannerStatus solved = rrt.solve(ptc);
      if (solved && !pdef->hasApproximateSolution()) {

        // estimate probability of success
        const ompl::base::PathPtr &ps = pdef->getSolutionPath();
        ompl::geometric::PathGeometric & gp(static_cast<ompl::geometric::PathGeometric&>(*ps));
        if (simplify) {
          psk->simplify(gp, ptc);
        }

        double prob = motionProbability(gp);
        visualizer->draw(prob,gp);

        // is this the best path?
        if (prob > bestProbability) {
          // yes, store it
          best = pdef;
          bestProbability = prob;
          bestStatus = solved;
        }

        // is probability above threshold?
        if (bestProbability > 0.75) {
          // yes, set as solution and return
          std::cout << "exact solution: probability = " << bestProbability << std::endl;
          pdef_->addSolutionPath(best->getSolutionPath());
          visualizer->clear();
          return ompl::base::PlannerStatus::EXACT_SOLUTION;
        }
      }
    }

    // take the best solution
    if (bestProbability > 0.0) {
      std::cout << "best solution: probability = " << bestProbability << std::endl;
      pdef_->addSolutionPath(best->getSolutionPath(), bestStatus == ompl::base::PlannerStatus::APPROXIMATE_SOLUTION);
      visualizer->clear();
      return bestStatus;
    }

    visualizer->clear();
    return ompl::base::PlannerStatus::TIMEOUT;
  }

};

#endif
